import os, shutil
import os.path as osp
import json
from tqdm import tqdm
from tool import filesystem, via_tool # export PYTHONPATH=$PYTHONPATH:`pwd`


def json_to_via(image_dir, 
                save_dir, 
                label_files,
                sub_dir="bdd100k", 
                via_name="via_region_data.json"):
    """
    bdd100k dataset tool
    将此数据集转换成via格式
    """
    total_files = filesystem.get_all_filepath(image_dir, [".jpg", "png"])
    file_json = dict()
    for f in total_files:
        file_json[osp.basename(f)] = f

    data_dict = list()
    for label_file in label_files:
        with open(label_file, "r") as rf:
            d = json.loads(rf.read())
            print(len(d))
            data_dict.extend(d)
    
    print(len(data_dict))
    category_set = set()
    def deal_one(idx, list_data, via_name="via_region_data.json"):
        cur_save_dir = osp.join(save_dir, sub_dir, str(idx))
        if not osp.exists(cur_save_dir):
            os.makedirs(cur_save_dir, exist_ok=True)
        total_json = dict()

        for data in tqdm(list_data):
            if not file_json.get(data["name"], None):continue

            file_path = file_json[data["name"]]
            file_size = osp.getsize(file_path)
            
            one_json = dict()
            one_json["filename"] = data["name"]
            one_json["size"] = file_size
            one_json["file_attributes"] = data["attributes"]
            regions = []
            for label in data.get("labels", [None]):
                if label is None:
                    print("label is None ", data["name"])
                    continue

                if not label.get("box2d", None): continue
                if label["category"] not in ["car", "bus", "truck", "pedestrian"] : continue
                
                region = dict()
                x = int(label["box2d"]["x1"])
                y = int(label["box2d"]["y1"])
                w = int(label["box2d"]["x2"] - x)
                h = int(label["box2d"]["y2"] - y)
                if w < 32 or h < 20:continue

                shape_attributes = {
                    "x": x,
                    "y": y,
                    "width": w,
                    "height": h,
                    "name": "rect"
                }
                region_attributes = label["attributes"]
                region_attributes["label"] = label["category"]
                category_set.add(label["category"])
                # region_attributes["manualShape"] = label["manualShape"]
                # region_attributes["manualAttributes"] = label["manualAttributes"]
                region_attributes["id"] = label["id"]
                
                region["shape_attributes"] = shape_attributes
                region["region_attributes"] = region_attributes
                regions.append(region)
            one_json["regions"] = regions
            total_json[data["name"] + str(file_size)] = one_json
            shutil.copy(file_path, cur_save_dir + os.sep + data["name"])
        
        with open(cur_save_dir + os.sep + via_name, "w") as wf:
            wf.write(json.dumps(total_json))

    # thread_nums = 12
    per_dir_imgs = 500
    for idx in range(len(data_dict) // per_dir_imgs):
        deal_one(idx, data_dict[idx*per_dir_imgs: (idx+1)* per_dir_imgs])
    print("category_set: ", category_set)

if __name__ == "__main__":
    # car truck bus
    total_labels = ['bicycle', 'traffic light', 'rider', 
        'pedestrian', 'other vehicle', 'trailer', 'motorcycle', 
        'truck', 'train', 'traffic sign', 'bus', 'other person', 
        'car']
        
    image_dir = "/home/xc/work/data/car/bdd100k/bdd100k/images/100k"
    save_dir = "/home/xc/work/data/car/train"
    label_files = [
        "/home/xc/work/data/car/bdd100k/bdd100k/labels/det_20/det_train.json",
        "/home/xc/work/data/car/bdd100k/bdd100k/labels/det_20/det_val.json"
    ]

    json_to_via(image_dir, save_dir, label_files)

